import os 

from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables import RunnableWithMessageHistory


os.environ["LANGSMITH_TRACING_V2"] = "true"
os.environ["LANGSMITH_API_KEY"] = "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"
os.environ["OPENAI_BASE_URL"] = "https://api.chatanywhere.tech/v1"
os.environ["OPENAI_API_KEY"] = "sk-pbXvhNj37SZ5SUBzC1Kx4LeXrsnT9EJNDL6mT2Lj2IbgohKa"

# 定义提示模板
prompt_template = ChatPromptTemplate.from_messages([
    ("system", "你是一个乐于助人的助手，用{language}尽你可能的回答所有的问题"),
    MessagesPlaceholder(variable_name="my_msg")
])
model = ChatOpenAI(model="gpt-4o-mini")
chain = prompt_template | model 

# 保存聊天
# 记录
store = {}
def get_session_history(session_id: str):
    if session_id not in store.keys():
        store[session_id] = ChatMessageHistory()
    return store[session_id]

do_message = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key="my_msg"
)


config = {
    "configurable": {
        "session_id": "1"
    }
}

if __name__ == "__main__":
    resp = do_message.invoke(
        {
            "language": "英文",
            "my_msg": [ HumanMessage(content="我是一个python开发工程师")]
        },
        config=config
    )
    print(resp.content)
    resp2 = do_message.invoke(
        {
            "language": "中文",
            "my_msg": [ HumanMessage(content="请给我推荐一本这方面专业的书籍")]
        },
        config=config
    )
    print(resp2.content)

    for resp in do_message.stream({"language": "英文", "my_msg": [HumanMessage(content="请给我讲一个笑话")]}, config=config):
        print(resp.content, end=" ")
